use crate::ast::{ASTNode, BinOp, SourceLocation};
use crate::errors::{CompilerError, ErrorKind};
use crate::type_system::TypeSystem;
use std::collections::HashMap;

// 类型信息
type TypeInfo = String;

// 类型环境
type TypeEnv = HashMap<String, TypeInfo>;

// 类型检查器结果
type TypeCheckResult<T> = Result<T, CompilerError>;

// 类型检查器
pub struct TypeChecker {
    env: TypeEnv,
    errors: Vec<CompilerError>,
    type_system: TypeSystem,
}

impl TypeChecker {
    // 创建类型检查器
    pub fn new() -> Self {
        TypeChecker {
            env: TypeEnv::new(),
            errors: Vec::new(),
            type_system: TypeSystem::new(),
        }
    }

    // 检查AST
    pub fn check(&mut self, node: &ASTNode) -> TypeCheckResult<()> {
        self.check_node(node)
    }

    // 检查节点
    fn check_node(&mut self, node: &ASTNode) -> TypeCheckResult<()> {
        match node {
            ASTNode::Program {
                preamble,
                body,
                epilogue,
                location: _,
            } => {
                for stmt in preamble {
                    self.check_node(stmt)?;
                }
                for stmt in body {
                    self.check_node(stmt)?;
                }
                for stmt in epilogue {
                    self.check_node(stmt)?;
                }
                Ok(())
            }
            ASTNode::VarDeclaration {
                name,
                type_name,
                value,
                location: _,
            } => {
                let value_type = self.infer_type(value)?;

                // 如果有类型标注，检查值类型是否与标注类型匹配
                if let Some(annotated_type) = type_name {
                    if !self.is_type_compatible(&value_type, &annotated_type) {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "变量 {} 类型不匹配: 标注类型 {}，值类型 {}",
                                name, annotated_type, value_type
                            ),
                        ));
                    }
                    // 使用标注的类型
                    self.env.insert(name.clone(), annotated_type.to_string());
                } else {
                    // 没有类型标注，使用推断的类型
                    self.env.insert(name.clone(), value_type);
                }
                Ok(())
            }
            ASTNode::ConstDeclaration {
                name,
                type_name,
                value,
                location: _,
            } => {
                let value_type = self.infer_type(value)?;

                // 如果有类型标注，检查值类型是否与标注类型匹配
                if let Some(annotated_type) = type_name {
                    if !self.is_type_compatible(&value_type, &annotated_type) {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "常量 {} 类型不匹配: 标注类型 {}，值类型 {}",
                                name, annotated_type, value_type
                            ),
                        ));
                    }
                    // 使用标注的类型
                    self.env.insert(name.clone(), annotated_type.to_string());
                } else {
                    // 没有类型标注，使用推断的类型
                    self.env.insert(name.clone(), value_type);
                }
                Ok(())
            }
            ASTNode::Assignment {
                name,
                value,
                location: _,
            } => {
                let value_type = self.infer_type(value)?;
                if let Some(var_type) = self.env.get(name) {
                    if var_type != &value_type {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "类型不匹配: 变量 {} 期望类型 {}, 得到类型 {}",
                                name, var_type, value_type
                            ),
                        ));
                    }
                } else {
                    self.env.insert(name.clone(), value_type);
                }
                Ok(())
            }
            ASTNode::BinaryOp {
                left,
                op,
                right,
                location,
            } => {
                let left_type = self.infer_type(left)?;
                let right_type = self.infer_type(right)?;

                // 检查二元操作的类型兼容性
                self.check_binary_op_compatibility(op, &left_type, &right_type, location)?;
                Ok(())
            }
            ASTNode::UnaryOp {
                op,
                operand,
                location,
            } => {
                let operand_type = self.infer_type(operand)?;

                // 检查一元操作的类型兼容性
                self.check_unary_op_compatibility(op, &operand_type, location)?;
                Ok(())
            }
            ASTNode::FunctionCall { name, args, .. } => {
                println!("处理函数调用: {}, 参数数量: {}", name, args.len());

                // 检查函数是否存在（在当前环境中或作为内置函数）
                let function_exists = self.env.contains_key(name) || self.is_builtin_function(name);
                println!(
                    "函数存在检查: {}, env中: {}, 内置函数: {}",
                    function_exists,
                    self.env.contains_key(name),
                    self.is_builtin_function(name)
                );

                if !function_exists {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::UndefinedIdentifier,
                        format!("未定义的函数: {}", name),
                    ));
                    // 继续检查参数，以发现更多潜在错误
                    for arg in args {
                        let arg_node = match arg {
                            crate::ast::Argument::Positional(node, _) => node,
                            crate::ast::Argument::Keyword(_, node, _) => node,
                        };
                        self.check_node(arg_node)?;
                    }
                    return Ok(());
                }

                // 对于内置函数，检查参数类型
                println!("是内置函数: {}", self.is_builtin_function(name));
                if self.is_builtin_function(name) {
                    // 检查参数数量
                    let expected_param_count = self.get_function_param_types("", name).len();
                    if args.len() != expected_param_count && expected_param_count > 0 {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "函数 {} 参数数量不匹配: 期望 {} 个参数，得到 {} 个参数",
                                name,
                                expected_param_count,
                                args.len()
                            ),
                        ));
                    }

                    // 检查参数类型
                    let expected_param_types = self.get_function_param_types("", name);
                    for (i, arg) in args.iter().enumerate() {
                        // 从Argument中提取ASTNode
                        let arg_node = match arg {
                            crate::ast::Argument::Positional(node, _) => node,
                            crate::ast::Argument::Keyword(_, node, _) => node,
                        };

                        self.check_node(arg_node)?;

                        if i < expected_param_types.len() {
                            let arg_type = self.infer_type(arg_node)?;
                            let expected_type = &expected_param_types[i];

                            println!(
                                "检查参数 {}: 期望 {}, 实际 {}",
                                i + 1,
                                expected_type,
                                arg_type
                            );

                            if !self.is_type_compatible(&arg_type, expected_type) {
                                self.errors.push(CompilerError::simple(
                                    ErrorKind::TypeMismatch,
                                    format!(
                                        "函数 {} 第 {} 个参数类型不匹配: 期望 {}，得到 {}",
                                        name,
                                        i + 1,
                                        expected_type,
                                        arg_type
                                    ),
                                ));
                                println!("检测到类型错误");
                            } else {
                                println!("类型兼容");
                            }
                        }
                    }

                    // 内置函数检查完成，返回
                    return Ok(());
                }

                // 如果不是内置函数，检查函数是否存在
                if !self.is_builtin_function(name) {
                    let func_type = if let Some(typ) = self.env.get(name) {
                        self.type_system.get_actual_type(typ)
                    } else {
                        // 函数未定义，添加错误
                        self.errors.push(CompilerError::simple(
                            ErrorKind::UndefinedFunction,
                            format!("未定义的函数: {}", name),
                        ));
                        return Ok(());
                    };

                    // 解析函数类型字符串，获取参数类型和返回类型
                    // 函数类型格式: (参数类型1,参数类型2,...)->返回类型
                    let (param_types_str, _return_type) = if let Some(pos) = func_type.find(")->") {
                        let param_part = &func_type[1..pos]; // 去掉开头的'('
                        let return_part = &func_type[pos + 2..]; // 跳过")->"
                        (param_part, return_part)
                    } else {
                        // 如果格式不正确，使用默认处理
                        self.get_function_param_types(&func_type, name);
                        return Ok(());
                    };

                    // 解析参数类型
                    let expected_param_types: Vec<String> = if param_types_str.is_empty() {
                        Vec::new()
                    } else {
                        param_types_str
                            .split(',')
                            .map(|s| s.trim().to_string())
                            .collect()
                    };

                    // 检查参数数量
                    if args.len() != expected_param_types.len() {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "函数 {} 参数数量不匹配: 期望 {} 个参数，得到 {} 个参数",
                                name,
                                expected_param_types.len(),
                                args.len()
                            ),
                        ));
                    }

                    // 检查每个参数的类型
                    for (i, arg) in args.iter().enumerate() {
                        // 从Argument中提取ASTNode
                        let arg_node = match arg {
                            crate::ast::Argument::Positional(node, _) => node,
                            crate::ast::Argument::Keyword(_, node, _) => node,
                        };

                        self.check_node(arg_node)?;

                        if i < expected_param_types.len() {
                            let arg_type = self.infer_type(arg_node)?;
                            let expected_type = &expected_param_types[i];

                            if !self.is_type_compatible(&arg_type, expected_type) {
                                self.errors.push(CompilerError::simple(
                                    ErrorKind::TypeMismatch,
                                    format!(
                                        "函数 {} 第 {} 个参数类型不匹配: 期望 {}，得到 {}",
                                        name,
                                        i + 1,
                                        expected_type,
                                        arg_type
                                    ),
                                ));
                            }
                        }
                    }
                }

                Ok(())
            }
            ASTNode::If {
                condition,
                then_block,
                else_if_conditions,
                else_if_blocks,
                else_block,
                location: _,
            } => {
                let cond_type = self.infer_type(condition)?;
                if cond_type != "真伪" {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::TypeMismatch,
                        "条件表达式必须是布尔类型".to_string(),
                    ));
                }

                for stmt in then_block {
                    self.check_node(stmt)?;
                }

                for (cond, block) in else_if_conditions.iter().zip(else_if_blocks.iter()) {
                    let cond_type = self.infer_type(cond)?;
                    if cond_type != "真伪" {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            "条件表达式必须是布尔类型".to_string(),
                        ));
                    }
                    for stmt in block {
                        self.check_node(stmt)?;
                    }
                }

                if let Some(else_stmt) = else_block {
                    for stmt in else_stmt {
                        self.check_node(stmt)?;
                    }
                }
                Ok(())
            }
            ASTNode::While {
                condition,
                body,
                location: _,
            } => {
                let cond_type = self.infer_type(condition)?;
                if cond_type != "真伪" {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::TypeMismatch,
                        "条件表达式必须是布尔类型".to_string(),
                    ));
                }

                for stmt in body {
                    self.check_node(stmt)?;
                }
                Ok(())
            }
            ASTNode::ForEach {
                variable,
                iterable,
                body,
                ..
            } => {
                let iter_type = self.infer_type(iterable)?;
                if iter_type != "集" {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::TypeMismatch,
                        "只能遍历列表类型".to_string(),
                    ));
                }

                // 为循环变量添加类型信息
                self.env.insert(variable.clone(), "数".to_string()); // 简化处理，实际应该根据列表元素类型推断
                for stmt in body {
                    self.check_node(stmt)?;
                }
                // 移除循环变量
                self.env.remove(variable);
                Ok(())
            }
            ASTNode::Return(expr, _) => {
                if let Some(e) = expr {
                    self.check_node(e)?;
                }
                Ok(())
            }
            ASTNode::FunctionDef {
                name,
                return_type,
                params,
                body,
                location: _,
            } => {
                // 创建新的作用域来存储函数参数
                let mut param_types = Vec::new();

                // 为函数参数添加类型信息
                for param in params {
                    // 如果参数有类型标注，使用标注的类型
                    // 否则，默认为"数"类型
                    let param_type = if let Some(ref type_name) = param.type_name {
                        type_name.clone()
                    } else {
                        "数".to_string() // 默认类型
                    };

                    // 存储参数类型，用于后续检查
                    param_types.push(param_type.clone());

                    // 将参数类型添加到环境中
                    self.env.insert(param.name.clone(), param_type);
                }

                // 检查函数体中的语句
                for stmt in body {
                    self.check_node(stmt)?;
                }

                // 如果有返回值类型标注，检查返回语句是否匹配
                if let Some(ret_type) = return_type {
                    // 简化处理：检查函数体中是否有返回语句
                    // 在实际实现中，应该更精确地分析所有可能的返回路径
                    for stmt in body {
                        if let ASTNode::Return(Some(expr), _) = stmt {
                            let expr_type = self.infer_type(expr)?;
                            if !self.is_type_compatible(&expr_type, &ret_type) {
                                self.errors.push(CompilerError::simple(
                                    ErrorKind::TypeMismatch,
                                    format!(
                                        "函数 {} 返回值类型不匹配: 期望 {}，得到 {}",
                                        name, ret_type, expr_type
                                    ),
                                ));
                            }
                        }
                    }
                }

                // 移除函数参数
                for param in params {
                    self.env.remove(&param.name);
                }

                // 将函数类型信息添加到环境中
                // 函数类型表示为: (参数类型1,参数类型2,...)->返回类型
                let params_str = param_types.join(",");
                let func_type = if let Some(ret_type) = return_type {
                    format!("({})->{}", params_str, ret_type)
                } else {
                    format!("({})->无", params_str)
                };

                self.env.insert(name.clone(), func_type);
                Ok(())
            }
            ASTNode::ClassDef {
                name: _,
                fields: _,
                methods: _,
                location: _,
            } => {
                // 简化处理，实际应该检查类的字段和方法
                Ok(())
            }
            ASTNode::MethodCall {
                object,
                method: _,
                args,
                location: _,
            } => {
                // 简化处理，实际应该检查方法调用
                self.check_node(object)?;
                for arg in args {
                    self.check_node(arg)?;
                }
                Ok(())
            }
            ASTNode::MemberAccess {
                object,
                field: _,
                location: _,
            } => {
                // 简化处理，实际应该检查成员访问
                self.check_node(object)?;
                Ok(())
            }
            ASTNode::MemberAssign {
                object,
                field: _,
                value,
                location: _,
            } => {
                // 简化处理，实际应该检查成员赋值
                self.check_node(object)?;
                self.check_node(value)?;
                Ok(())
            }
            ASTNode::ArrayAccess {
                array,
                index,
                location: _,
            } => {
                let array_type = self.infer_type(array)?;
                if array_type != "集" && array_type != "文" {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::TypeMismatch,
                        "只能对列表或字符串进行索引访问".to_string(),
                    ));
                }

                let index_type = self.infer_type(index)?;
                if index_type != "数" {
                    self.errors.push(CompilerError::simple(
                        ErrorKind::TypeMismatch,
                        "索引必须是数字类型".to_string(),
                    ));
                }
                Ok(())
            }
            ASTNode::Array(elements, _) => {
                for elem in elements {
                    self.check_node(elem)?;
                }
                Ok(())
            }
            ASTNode::ArrayPush {
                array,
                value,
                location: _,
            } => {
                self.check_node(array)?;
                self.check_node(value)?;
                Ok(())
            }
            // 其他节点类型
            _ => Ok(()),
        }
    }

    // 推断类型
    fn infer_type(&mut self, node: &ASTNode) -> TypeCheckResult<TypeInfo> {
        match node {
            ASTNode::Number(_, _) => Ok("数".to_string()),
            ASTNode::String(_, _) => Ok("文".to_string()),
            ASTNode::Boolean(_, _) => Ok("真伪".to_string()),
            ASTNode::Null(_) => Ok("无".to_string()),
            ASTNode::Array(_, _) => Ok("集".to_string()),
            ASTNode::Identifier(name, _) => {
                if let Some(typ) = self.env.get(name) {
                    // 获取实际类型，处理类型别名
                    Ok(self.type_system.get_actual_type(typ))
                } else {
                    Err(CompilerError::simple(
                        ErrorKind::UndefinedIdentifier,
                        format!("未定义的标识符: {}", name),
                    ))
                }
            }
            ASTNode::BinaryOp {
                left, op, right, ..
            } => {
                let left_type = self.infer_type(left)?;
                let right_type = self.infer_type(right)?;

                // 根据操作推断结果类型
                match (op, left_type.as_str(), right_type.as_str()) {
                    (
                        BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Mod,
                        "数",
                        "数",
                    ) => Ok("数".to_string()),
                    (
                        BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge | BinOp::Eq | BinOp::Ne,
                        _,
                        _,
                    ) => Ok("真伪".to_string()),
                    (BinOp::And | BinOp::Or, "真伪", "真伪") => Ok("真伪".to_string()),
                    (BinOp::Add, "文", "文") => Ok("文".to_string()),
                    _ => {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!(
                                "不支持的操作数类型组合: {} {} {}",
                                left_type, op, right_type
                            ),
                        ));
                        Ok("无".to_string())
                    }
                }
            }
            ASTNode::UnaryOp { op, operand, .. } => {
                let operand_type = self.infer_type(operand)?;

                match (op.as_str(), operand_type.as_str()) {
                    ("-", "数") => Ok("数".to_string()),
                    ("!", "真伪") => Ok("真伪".to_string()),
                    _ => {
                        self.errors.push(CompilerError::simple(
                            ErrorKind::TypeMismatch,
                            format!("不支持的一元操作数类型: {} {}", op, operand_type),
                        ));
                        Ok("无".to_string())
                    }
                }
            }
            ASTNode::FunctionCall { name, args: _, .. } => {
                // 检查函数是否存在
                let function_exists = self.env.contains_key(name) || self.is_builtin_function(name);

                if !function_exists {
                    return Err(CompilerError::simple(
                        ErrorKind::UndefinedIdentifier,
                        format!("未定义的函数: {}", name),
                    ));
                }

                // 获取函数返回类型
                if let Some(func_type) = self.env.get(name) {
                    // 从函数类型中提取返回类型
                    if func_type.starts_with("函数(") && func_type.contains(") -> ") {
                        let return_part = func_type.split(") -> ").last().unwrap();
                        Ok(return_part.to_string())
                    } else {
                        Ok("无".to_string())
                    }
                } else if self.is_builtin_function(name) {
                    // 内置函数返回类型
                    match name.as_str() {
                        "书" | "书行" | "推" | "弹" => Ok("无".to_string()),
                        "包含" | "空" => Ok("真伪".to_string()),
                        "长" => Ok("数".to_string()),
                        "取" => Ok("无".to_string()), // 简化处理
                        _ => Ok("无".to_string()),
                    }
                } else {
                    Ok("无".to_string())
                }
            }
            // 其他节点类型，简化处理
            _ => Ok("无".to_string()),
        }
    }

    // 检查二元操作兼容性
    fn check_binary_op_compatibility(
        &mut self,
        op: &BinOp,
        left_type: &TypeInfo,
        right_type: &TypeInfo,
        _location: &Option<SourceLocation>,
    ) -> TypeCheckResult<()> {
        let compatible = match (op, left_type.as_str(), right_type.as_str()) {
            (BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Mod, "数", "数") => true,
            (BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge, "数", "数") => true,
            (BinOp::Eq | BinOp::Ne, _, _) => true,
            (BinOp::And | BinOp::Or, "真伪", "真伪") => true,
            (BinOp::Add, "文", "文") => true,
            _ => false,
        };

        if !compatible {
            self.errors.push(CompilerError::simple(
                ErrorKind::TypeMismatch,
                format!(
                    "不支持的操作数类型组合: {} {} {}",
                    left_type, op, right_type
                ),
            ));
        }

        Ok(())
    }

    // 检查一元操作兼容性
    fn check_unary_op_compatibility(
        &mut self,
        op: &String,
        operand_type: &TypeInfo,
        _location: &Option<SourceLocation>,
    ) -> TypeCheckResult<()> {
        let compatible = match (op.as_str(), operand_type.as_str()) {
            ("-", "数") => true,
            ("!", "真伪") => true,
            _ => false,
        };

        if !compatible {
            self.errors.push(CompilerError::simple(
                ErrorKind::TypeMismatch,
                format!("不支持的一元操作数类型: {} {}", op, operand_type),
            ));
        }

        Ok(())
    }

    // 获取错误信息
    pub fn errors(&self) -> &[CompilerError] {
        &self.errors
    }

    // 是否有错误
    pub fn has_errors(&self) -> bool {
        !self.errors.is_empty()
    }

    // 检查是否为内置函数
    fn is_builtin_function(&self, name: &str) -> bool {
        matches!(
            name,
            "书" | "书行"
                | "使"
                | "入"
                | "出"
                | "加"
                | "减"
                | "乘"
                | "除"
                | "模"
                | "等"
                | "不等"
                | "小于"
                | "大于"
                | "小于等于"
                | "大于等于"
                | "且"
                | "或"
                | "非"
                | "包含"
                | "空"
                | "型"
                | "推"
                | "弹"
                | "取"
                | "定"
                | "随"
                | "时"
                | "数"
                | "文"
                | "长"
        )
    }

    // 获取内置函数类型
    #[allow(unused)]
    fn get_builtin_function_type(&self, name: &str) -> String {
        match name {
            // 输出函数
            "书" | "书行" => "函数(文) -> 无".to_string(),

            // 数学函数
            "加" | "减" | "乘" | "除" | "模" => "函数(数, 数) -> 数".to_string(),

            // 比较函数
            "等" | "不等" | "小于" | "大于" | "小于等于" | "大于等于" => {
                "函数(数, 数) -> 真伪".to_string()
            }

            // 逻辑函数
            "且" | "或" => "函数(真伪, 真伪) -> 真伪".to_string(),
            "非" => "函数(真伪) -> 真伪".to_string(),

            // 字符串函数
            "长" => "函数(文) -> 数".to_string(),
            "包含" => "函数(文, 文) -> 真伪".to_string(),

            // 数组函数
            "推" => "函数(集, 无) -> 无".to_string(),
            "弹" => "函数(集) -> 无".to_string(),
            "取" => "函数(集, 数) -> 无".to_string(),

            // 类型转换
            "数" => "函数(文) -> 数".to_string(),
            "文" => "函数(数) -> 文".to_string(),

            // 其他
            _ => "函数(...) -> 无".to_string(),
        }
    }

    // 获取函数参数类型列表
    fn get_function_param_types(&self, _func_type: &str, func_name: &str) -> Vec<String> {
        // 简化处理：从内置函数类型中提取参数类型
        if self.is_builtin_function(func_name) {
            match func_name {
                "书" | "书行" => vec!["文".to_string()], // 书和书行函数都需要一个字符串参数
                "非" => vec!["真伪".to_string()],        // 非函数需要一个布尔参数
                "弹" => vec!["集".to_string()],          // 弹函数需要一个集合参数
                "长" => vec!["文".to_string()],          // 长函数需要一个字符串参数

                "包含" => vec!["文".to_string(), "文".to_string()],
                "推" => vec!["集".to_string(), "无".to_string()],
                "取" => vec!["集".to_string(), "数".to_string()],

                _ => vec![],
            }
        } else {
            // 对于用户定义的函数，简化处理
            vec!["文".to_string()] // 假设所有未定义函数都需要一个字符串参数
        }
    }

    // 检查类型兼容性
    fn is_type_compatible(&self, actual_type: &str, expected_type: &str) -> bool {
        // 完全匹配
        if actual_type == expected_type {
            return true;
        }

        // 处理数字类型（整数和浮点数都视为"数"）
        if expected_type == "数" && (actual_type == "整数" || actual_type == "浮点") {
            return true;
        }

        // 处理字符串类型
        if expected_type == "文" && actual_type == "字符串" {
            return true;
        }

        // 处理布尔类型
        if expected_type == "真伪" && actual_type == "布尔" {
            return true;
        }

        // 处理集合类型（集和表）
        if expected_type == "集" && actual_type == "表" {
            return true;
        }

        if expected_type == "表" && actual_type == "集" {
            return true;
        }

        // 检查类型别名
        let actual = self.type_system.get_actual_type(actual_type);
        let expected = self.type_system.get_actual_type(expected_type);

        actual == expected
    }
}
